import zipfile
import os

def zip_epub(source_dir, book_name=None, author_name=None):
    # 检查source_dir是否存在
    if not os.path.exists(source_dir) or not os.path.isdir(source_dir):
        print(f"错误: EPUB目录 '{source_dir}' 不存在")
        return False

    # 获取输出目录
    output_dir = os.path.dirname(source_dir)
    
    # 生成文件名
    if book_name and author_name:
        epub_name = os.path.join(output_dir, f"《{book_name}》作者：{author_name}.epub")
    else:
        epub_name = os.path.join(output_dir, f"{os.path.basename(source_dir)}.epub")
    
    print(f"开始生成EPUB文件: '{epub_name}' ")
    
    try:
        # 创建epub文件
        with zipfile.ZipFile(epub_name, 'w') as epub:
            # 首先添加mimetype文件，不压缩
            mimetype_path = os.path.join(source_dir, 'mimetype')
            if os.path.exists(mimetype_path):
                epub.write(mimetype_path, 'mimetype', compress_type=zipfile.ZIP_STORED)
            
            # 添加其他所有文件，使用最高压缩级别
            for root, dirs, files in os.walk(source_dir):
                for file in files:
                    if file != 'mimetype':
                        file_path = os.path.join(root, file)
                        arcname = os.path.relpath(file_path, source_dir)
                        
                        # 对不同类型的文件使用不同的压缩方式
                        ext = os.path.splitext(file)[1].lower()
                        # 已压缩的文件格式使用STORED
                        if ext in ['.jpg', '.jpeg', '.png', '.gif', '.mp3', '.mp4']:
                            compress_type = zipfile.ZIP_STORED
                        # 其他文件使用最高压缩级别
                        else:
                            compress_type = zipfile.ZIP_DEFLATED
                        
                        epub.write(file_path, arcname, compress_type=compress_type,
                                compresslevel=9)  # 使用最高压缩级别
        print(f"书籍文件成功生成：{epub_name} ")
        return True
    except Exception as e:
        print(f"生成EPUB文件时出错: {str(e)}")
        return False